package com.baomidou.mybatisplus.ext.util;

import com.sun.source.tree.Tree;
import com.sun.source.util.TreePath;
import com.sun.source.util.Trees;
import com.sun.tools.javac.code.Flags;
import com.sun.tools.javac.model.JavacElements;
import com.sun.tools.javac.tree.JCTree;
import com.sun.tools.javac.tree.TreeMaker;
import com.sun.tools.javac.util.List;
import com.sun.tools.javac.util.Names;

import javax.lang.model.element.Element;
import java.util.ArrayList;

/**
 * Java抽象语法树AST 工具类
 *
 * @author wanglei
 */
public class AstUtil {


    /**
     * 定义一个变量并且new 对象(无构造参数请传：com.sun.tools.javac.util.List.nil())
     *
     * @param maker
     * @param names
     * @param classAllName
     * @param varName
     * @return
     */
    public static JCTree.JCVariableDecl addField(TreeMaker maker, JavacElements elementUtils, Names names, String classAllName, String varName,List<JCTree.JCExpression> constructorParams) {
        return maker.VarDef(maker.Modifiers(Flags.PRIVATE), names.fromString(varName), memberAccess(maker, elementUtils, classAllName),
                newObject(maker, names, elementUtils,classAllName,constructorParams));
    }

    /**
     * 调用一个方法
     * @param maker
     * @param elementUtils
     * @param method 方法名
     * @param types 类型
     * @param paramsList 参数列表
     * @return
     */
    public static JCTree.JCMethodInvocation callMethod(TreeMaker maker, JavacElements elementUtils, String method, String[] types, List<JCTree.JCExpression> paramsList) {
        List<JCTree.JCExpression> paramsTypeList = List.nil();
        for (String type : types) {
            paramsTypeList = paramsTypeList.append(memberAccess(maker, elementUtils, type));
        }
       return  maker.Apply(paramsTypeList, memberAccess(maker,elementUtils,method), paramsList);
    }

    /**
     * 创建多级访问
     *
     * @param components
     * @return
     */
    public static JCTree.JCExpression memberAccess(TreeMaker treeMaker, JavacElements elementUtils, String components) {
        String[] componentArray = components.split("\\.");
        JCTree.JCExpression expr = treeMaker.Ident(elementUtils.getName(componentArray[0]));
        for (int i = 1; i < componentArray.length; i++) {
            expr = treeMaker.Select(expr, elementUtils.getName(componentArray[i]));
        }
        return expr;
    }


    /**
     * 初始化一个对象
     *
     * @param maker
     * @param names
     * @param classAllName
     * @return
     */
    public static JCTree.JCNewClass newObject(TreeMaker maker, Names names,JavacElements elementUtils, String classAllName,List<JCTree.JCExpression> constructorParams) {
        return maker.NewClass(null,
                com.sun.tools.javac.util.List.nil(),
                memberAccess(maker, elementUtils, classAllName),
                constructorParams,
                null);
    }

    /**
     * 导入包
     *
     * @param element
     * @param maker
     * @param names
     * @param tree
     * @param packageName 报名
     * @param className   类名
     */
    public static void addImportInfo(Element element, TreeMaker maker, Names names, Trees tree, String packageName, String className) {
        TreePath treePath = tree.getPath(element);
        Tree leaf = treePath.getLeaf();
        if (treePath.getCompilationUnit() instanceof JCTree.JCCompilationUnit && leaf instanceof JCTree) {
            JCTree.JCCompilationUnit jccu = (JCTree.JCCompilationUnit) treePath.getCompilationUnit();

            for (JCTree jcTree : jccu.getImports()) {
                if (jcTree != null && jcTree instanceof JCTree.JCImport) {
                    JCTree.JCImport jcImport = (JCTree.JCImport) jcTree;
                    if (jcImport.qualid != null && jcImport.qualid instanceof JCTree.JCFieldAccess) {
                        JCTree.JCFieldAccess jcFieldAccess = (JCTree.JCFieldAccess) jcImport.qualid;
                        try {
                            if (packageName.equals(jcFieldAccess.selected.toString()) && className.equals(jcFieldAccess.name.toString())) {
                                return;
                            }
                        } catch (NullPointerException e) {
                            e.printStackTrace();
                        }
                    }
                }
            }
            java.util.List<JCTree> trees = new ArrayList<>();
            trees.addAll(jccu.defs);
            JCTree.JCIdent ident = maker.Ident(names.fromString(packageName));
            JCTree.JCImport jcImport = maker.Import(maker.Select(
                    ident, names.fromString(className)), false);
            if (!trees.contains(jcImport)) {
                trees.add(0, jcImport);
            }
            jccu.defs = List.from(trees);
        }
    }
}
